#encoding=utf-8
import copy
import socket
from utils.common_import import *
from config.cluster_config_new import cluster_client
from common.kubenetes.k8sNew import KubeClient
from kubernetes.client import V1Pod
from typing import List

gpu_device_name = cluster_client.cluster_config["device"]["gpu"]["resource_name"]

def get_non_terminated_pods(kube_client: KubeClient=cluster_client.kube_client) -> List[V1Pod]:
    return kube_client.corev1_client.list_pod_for_all_namespaces(
        field_selector = r'status.phase!=Succeeded,status.phase!=Failed,spec.nodeName!=%22%22').items

#获取节点资源使用信息
def statistic_node_resource(kube_client: KubeClient=cluster_client.kube_client):

    default_resouce = {"gpu":
                           {"total":0,
                            "left": 0,
                            'unavailable': ''
                            },
                       "mem":
                           {"total":0,
                            "left": 0},
                       "cpu":
                           {"total":0,
                            "left": 0
                            },
                       "machineType": "default"
                       }

    pods = get_non_terminated_pods(kube_client)
    request_resource = {}
    for pod in pods:
        host_name = pod.spec.node_name
        request_resource.setdefault(host_name, {'cpu': 0, 'memory': 0, gpu_device_name: 0})
        for container in pod.spec.containers:
            if not container.resources.requests:
                continue
            for resource, value in container.resources.requests.items():
                if resource in ['cpu', 'memory', gpu_device_name]:
                    if value[-2:] in ["Ki", "Mi", "Gi"]:
                        if value[-2:] == "Ki":
                            request_resource[host_name][resource] += float(value[:-2])/1024
                        elif value[-2:] == "Mi":
                            request_resource[host_name][resource] += float(value[:-2])
                        else:
                            request_resource[host_name][resource] += float(value[:-2])*1024
                    elif value[-1] == 'm':
                        request_resource[host_name][resource] += float(value[:-1])/1000
                    else:
                        request_resource[host_name][resource] += float(value)

    nodes_resource = {}
    nodes = kube_client.get_nodes()
    for node in nodes.items:
        host_ip = ''
        host_name = ''
        for address in node.status.addresses:
            if address.type == 'InternalIP':
                host_ip = address.address
            if address.type == 'Hostname':
                host_name = address.address
        if not host_ip or not host_name:
            continue

        node_total_memory = node.status.capacity['memory']
        if node_total_memory[-2:] in ["Ki", "Mi", "Gi"]:
            if node_total_memory[-2:] == "Ki":
                node_total_memory = float(node_total_memory[:-2]) / 1024
            elif node_total_memory[-2:] == "Mi":
                node_total_memory = float(node_total_memory[:-2])
            else:
                node_total_memory = float(node_total_memory[:-2]) * 1024

            node_total_memory = float("{:.2f}".format(node_total_memory))

        node_total_cpu = float(node.status.capacity['cpu'])
        node_total_gpu = float(node.status.capacity.get(gpu_device_name, 0))
        if node_total_cpu <= 0 or node_total_memory <= 0:
            continue

        nodes_resource.setdefault(host_ip, copy.deepcopy(default_resouce))
        nodes_resource[host_ip]["gpu"]["left"] = node_total_gpu
        if host_name in request_resource:
            nodes_resource[host_ip]["gpu"]["left"] -= request_resource[host_name].get(gpu_device_name, 0)
        nodes_resource[host_ip]["gpu"]["total"] += node_total_gpu

        nodes_resource[host_ip]["mem"]["left"] = node_total_memory
        if host_name in request_resource:
            nodes_resource[host_ip]["mem"]["left"] -= request_resource[host_name].get('memory', 0)
            nodes_resource[host_ip]["mem"]["left"] = float("{:.2f}".format(nodes_resource[host_ip]["mem"]["left"]))
        nodes_resource[host_ip]["mem"]["total"] += node_total_memory

        nodes_resource[host_ip]["cpu"]["left"] += node_total_cpu
        if host_name in request_resource:
            nodes_resource[host_ip]["cpu"]["left"] -= request_resource[host_name].get('cpu', 0)
            nodes_resource[host_ip]["cpu"]["left"] = float("{:.2f}".format(nodes_resource[host_ip]["cpu"]["left"]))
        nodes_resource[host_ip]["cpu"]["total"] += node_total_cpu

        nodes_resource[host_ip]["machineType"] = node.metadata.labels.get("resourceType", "default")

    logger.info("[statistic_node_resource]: get resource by gpu type: {}".format(nodes_resource))
    return nodes_resource


#获取节点gpu类型以及节点使用信息
def get_node_info(kube_client: KubeClient):
    nodes_resource = statistic_node_resource(kube_client)
    logger.info("nodes_resource {}".format(nodes_resource))

    #按数值大小排序
    result = sorted(nodes_resource.items(), key= lambda item: socket.inet_aton(item[0]))
    nodes_resource_list = []
    for item in result:
        #转为数组，添加节点id
        item[1]['nodeId'] = item[0]
        nodes_resource_list.append(item[1])

    logger.info(nodes_resource_list)

    return nodes_resource_list

#根据集群名称获取节点资源信息
def get_nodes_resource_by_cluster_name():

    return  get_node_info(cluster_client.kube_client)

if __name__ == '__main__':
    logger.info(get_nodes_resource_by_cluster_name())

